from langchain.agents import Tool
import requests
import json
from .travel_assistant_agent import travel_assistant
from .agent_selector import get_access_token

def travel_creator_with_text(user_input):

    generated_text = travel_assistant(user_input)

    # 使用生成的文本内容直接作为大模型的输入
    system_text = (
        "You are a prompt creator specializing in generating high-quality, photorealistic prompts for Stable-Diffusion-XL. Your goal is to transform input text into vivid, highly detailed descriptions that capture travelers enjoying iconic landmarks and landscapes in a realistic, high-definition style. For instance, if the setting is a historic city, depict travelers happily exploring its famous sites. Focus on authentic architecture, scenic elements unique to the area, natural lighting, vibrant colors, and true-to-life details. Avoid any anime, cartoon, or stylized effects, ensuring that the image is as realistic as possible. Please respond in English only, with a maximum word limit of 60."
    )
    full_input = system_text + generated_text

    url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token={get_access_token()}"
    payload = json.dumps({
        "messages": [{"role": "user", "content": full_input}],
        "disable_search": False,
        "system": system_text
    })
    headers = {'Content-Type': 'application/json'}
    response = requests.post(url, headers=headers, data=payload)

    # 输出生成的绘图提示词
    prompt_result = response.json().get("result", "无返回结果")
    return prompt_result

# 定义 travel_draw_tool，将 travel_creator_with_text 封装为 Tool
travel_draw_tool = Tool(
    name="Drawing Prompt Generator",
    func=travel_creator_with_text,
    description="直接根据生成的内容生成绘图提示词。"
)

# 定义 generate_image 函数，用于基于提示词生成图像
def generate_image(prompt):
    access_token = get_access_token()
    url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/text2image/sd_xl?access_token={access_token}"

    payload = json.dumps({
        "prompt": prompt,
        "negative_prompt": "卡通,动漫,漫画",
        "size": "768x768",
        "n": 1,
        "steps": 40,
        "sampler_index": "DPM2"
    })

    headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}

    # 发送请求并打印调试信息
    response = requests.post(url, headers=headers, data=payload)
    print("图像生成 API 响应状态码:", response.status_code)
    print("图像生成 API 响应内容:", response.text)

    if response.status_code == 200:
        data = response.json()
        if 'data' in data and len(data['data']) > 0:
            return data['data'][0].get('b64_image')
        else:
            print("图像生成响应中无图像数据")
            return None
    else:
        print("图像生成请求失败:", response.text)
        return None
